import sqlite3
import json
import os
import threading
import time
from contextlib import contextmanager


def retry_on_locked(max_retries=100, delay=0.01):
    """
    重试装饰器,用于处理SQLite锁占用错误
    :param max_retries: 最大重试次数
    :param delay: 重试间隔时间（秒）
    """
    def decorator(func):
        def wrapper(*args, **kwargs):
            retries = 0
            last_error = None
            while retries < max_retries:
                try:
                    return func(*args, **kwargs)
                except sqlite3.OperationalError as e:
                    if "locked" in str(e).lower() or "database is locked" in str(e).lower():
                        last_error = e
                        retries += 1
                        time.sleep(delay)
                        continue
                    raise  # 如果不是锁占用错误，直接抛出
            # 如果重试次数用完，抛出最后的错误
            raise last_error if last_error else sqlite3.OperationalError("Database locked after retries")
        return wrapper
    return decorator


class SQLiteDict:
    def __init__(self, db_path, table_name='kv_store'):
        """
        初始化SQLite数据库实例
        :param db_path: 数据库文件路径
        :param table_name: 表名(默认为'kv_store')
        """
        self.db_path = db_path
        self.table_name = table_name
        self._local = threading.local()
        self._init_db()

    def _get_connection(self):
        """获取线程本地连接"""
        if not hasattr(self._local, 'connection'):
            self._local.connection = sqlite3.connect(self.db_path)
            self._local.connection.row_factory = sqlite3.Row
        return self._local.connection

    @retry_on_locked()
    def _init_db(self):
        """初始化数据库表"""
        conn = self._get_connection()
        cursor = conn.cursor()
        cursor.execute(f'''
            CREATE TABLE IF NOT EXISTS {self.table_name} (
                key TEXT PRIMARY KEY,
                value TEXT NOT NULL
            )
        ''')
        conn.commit()

    @retry_on_locked()
    def __getitem__(self, key):
        if not isinstance(key, str):
            raise TypeError("键必须是字符串类型")

        conn = self._get_connection()
        cursor = conn.cursor()
        cursor.execute(f'SELECT value FROM {self.table_name} WHERE key = ?', (key,))
        row = cursor.fetchone()
        if row is None:
            raise KeyError(key)
        return row['value']

    @retry_on_locked()
    def __setitem__(self, key, value):
        if not isinstance(key, str):
            raise TypeError("键必须是字符串类型")
        if not isinstance(value, str):
            raise TypeError("值必须是字符串类型")

        conn = self._get_connection()
        cursor = conn.cursor()
        cursor.execute(f'''
            INSERT OR REPLACE INTO {self.table_name} (key, value) VALUES (?, ?)
        ''', (key, value))
        conn.commit()

    @retry_on_locked()
    def __delitem__(self, key):
        if not isinstance(key, str):
            raise TypeError("键必须是字符串类型")

        conn = self._get_connection()
        cursor = conn.cursor()
        cursor.execute(f'DELETE FROM {self.table_name} WHERE key = ?', (key,))
        if cursor.rowcount == 0:
            raise KeyError(key)
        conn.commit()

    @retry_on_locked()
    def __contains__(self, key):
        if not isinstance(key, str):
            return False

        conn = self._get_connection()
        cursor = conn.cursor()
        cursor.execute(f'SELECT 1 FROM {self.table_name} WHERE key = ?', (key,))
        return cursor.fetchone() is not None

    def __iter__(self):
        for key in self.keys():
            yield key

    @retry_on_locked()
    def keys(self):
        """返回键的生成器"""
        conn = self._get_connection()
        cursor = conn.cursor()
        cursor.execute(f'SELECT key FROM {self.table_name}')
        for row in cursor:
            yield row['key']

    @retry_on_locked()
    def values(self):
        """返回值的生成器"""
        conn = self._get_connection()
        cursor = conn.cursor()
        cursor.execute(f'SELECT value FROM {self.table_name}')
        for row in cursor:
            yield row['value']

    @retry_on_locked()
    def items(self):
        """返回键值对的生成器"""
        conn = self._get_connection()
        cursor = conn.cursor()
        cursor.execute(f'SELECT key, value FROM {self.table_name}')
        for row in cursor:
            yield (row['key'], row['value'])

    def get(self, key, default=None):
        """获取键对应的值，如果不存在则返回默认值"""
        try:
            return self[key]
        except KeyError:
            return default

    def pop(self, key, default=None):
        """删除键并返回其值，如果不存在则返回默认值"""
        try:
            value = self[key]
            del self[key]
            return value
        except KeyError:
            return default

    @retry_on_locked()
    def clear(self):
        """清空所有数据"""
        conn = self._get_connection()
        cursor = conn.cursor()
        cursor.execute(f'DELETE FROM {self.table_name}')
        conn.commit()

    @retry_on_locked()
    def __len__(self):
        """返回键值对数量"""
        conn = self._get_connection()
        cursor = conn.cursor()
        cursor.execute(f'SELECT COUNT(*) FROM {self.table_name}')
        return cursor.fetchone()[0]

    def close(self):
        """关闭数据库连接"""
        if hasattr(self._local, 'connection'):
            self._local.connection.close()
            delattr(self._local, 'connection')

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
